import time
import multiprocessing
import ffmpeg


class stream_pusher(object):
    def __init__(self, rtmp_url=None, raw_frame_q=None, fps=25, width=1920, height=1080):  # 类实例化的时候传入rtmp地址和帧传入队列
        self.rtmp_url = rtmp_url
        self.raw_frame_q = raw_frame_q
        self.fps = fps  # 设置帧速率
        # 设置分辨率
        self.width = int(width)  # 宽
        self.height = int(height)  # 高

    # 对获取的帧做一些画面处理的方法，返回完成处理的帧。
    def __frame_handle__(self, raw_frame, text, shape1, shape2):
        # 帧用cv2进行一些处理，比如写上文本，画矩形等
        return (raw_frame)

    def push_frame_to_rtmp(self):
        try:
            width, height = self.width, self.height
            while True:
                if not self.raw_frame_q.empty():  # 如果输入管道不为空
                    # 把帧和相关信息从输入队列中取出
                    raw_frame, text, shape1, shape2 = self.raw_frame_q.get()
                    # 将帧数据推送到 RTMP 服务器
                    print('starting...')
                    ffmpeg_path = r'D:\通用软件\ffmpeg\bin\ffmpeg.exe'
                    ffmpeg.input('pipe:0',
                                 format='rawvideo',
                                 pix_fmt='bgr24',
                                 s=f'{width}x{height}').output(self.rtmp_url,
                                                               format='flv',
                                                               codec='libx264',
                                                               preset="veryfast",
                                                               tune="zerolatency").run(input=raw_frame.tobytes(),
                                                                                       cmd=ffmpeg_path)
                else:
                    time.sleep(0.01)

        except ffmpeg.Error as e:
            print("Error occurred while pushing frame to RTMP:")
            print(e)

    # 启动运行
    def run(self):
        # 定义一个子进程
        push_frame_p = multiprocessing.Process(target=self.push_frame_to_rtmp, args=())
        push_frame_p.daemon = True  # 把子进程设置为daemon方式
        push_frame_p.start()  # 运行子进程